package main

import (
	"code.google.com/p/bouncer/server/utils"
	"errors"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"sync"
	"time"
)

var maxConns = flag.Int("max-conns", 500,
	"Maximum number of connected clients")
var maxBytesPerConn = flag.Int64("max-bytes-per-conn", 1024*500,
	"Maximum number of bytes in a given direction per proxied connection")
var maxConnsPerIp = flag.Int("max-conns-per-ip", 10,
	"Maximum number of connected clients per IP")
var rendezvousTimeout = flag.Int64("initial-timeout-secs", 5,
	"How much time to wait for a peer before timing out a client")
var inactivityTimeout = flag.Int64("inactivity-timeout-secs", 10,
	"How much time without network events before killing a proxy session")

func ipToKey(ip net.IP) string {
	if ip.To4() == nil {
		ret := make([]byte, 16)
		copy(ret[:8], ip[:8])
		return string(ret)
	}
	return string(ip.To16())
}

type Accounting struct {
	sync.Mutex
	conns      int
	connsPerIp map[string]int
}

func (a *Accounting) TryAdd(ip net.IP) bool {
	key := ipToKey(ip)
	a.Lock()
	defer a.Unlock()
	if a.conns >= *maxConns {
		return false
	}
	perIp, ok := a.connsPerIp[key]
	if ok && perIp >= *maxConnsPerIp {
		return false
	}
	a.connsPerIp[key] = perIp + 1
	a.conns++
	return true
}

func (a *Accounting) Remove(ip net.IP) {
	key := ipToKey(ip)
	a.Lock()
	defer a.Unlock()
	a.conns--
	if a.connsPerIp[key] == 1 {
		delete(a.connsPerIp, key)
	} else {
		a.connsPerIp[key]--
	}
}

// Connection wrapper that limits the total number of bytes read, and
// releases a connection resource when closed.
type AccountingConn struct {
	net.Conn
	read   int64
	closed bool
}

func (a *AccountingConn) Read(b []byte) (int, error) {
	if a.read >= *maxBytesPerConn {
		return 0, errors.New(
			fmt.Sprintf("Byte quota exceeded by %s", a.Conn.RemoteAddr()))
	}
	n, err := a.Conn.Read(b)
	a.read += int64(n)
	return n, err
}

func (a *AccountingConn) Close() error {
	if !a.closed {
		limits.Remove(a.Conn.RemoteAddr().(*net.TCPAddr).IP)
		a.closed = true
	}
	return a.Conn.Close()
}

func hijack(w http.ResponseWriter) (net.Conn, error) {
	// Stop the Go library from telling the client to expect chunked
	// content encoding.
	w.Header().Add("content-length", "0")
	w.Header().Add("upgrade", "TCP")
	w.WriteHeader(101)
	if h, ok := w.(http.Hijacker); ok {
		conn, buf, err := h.Hijack()
		if err != nil {
			return nil, err
		}
		// The buffer does not contain the final \r\n for the HTTP
		// headers, so we need to add that ourselves. With that extra
		// newline, the HTTP headers are complete and the client will
		// be switching protocols.
		buf.Writer.WriteString("\r\n")
		buf.Flush()
		return &AccountingConn{Conn: conn}, nil
	}
	panic("HTTP hijacking not supported")
}

func handle(w http.ResponseWriter, r *http.Request) {
	if r.Header.Get("upgrade") != "TCP" {
		w.WriteHeader(400)
		w.Write([]byte("Missing 'Upgrade: TCP' header"))
		return
	}

	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		w.WriteHeader(400)
		return
	}
	ip := net.ParseIP(host)
	if ip == nil {
		w.WriteHeader(400)
		return
	}
	if !limits.TryAdd(ip) {
		w.WriteHeader(503)
		w.Write([]byte("Server full, try again later"))
		return
	}

	err = match.Add(r.URL.Path, func(first bool) (net.Conn, error) {
		var role string
		if first {
			role = "first"
		} else {
			role = "second"
		}
		w.Header().Add("x-bouncer-order", role)
		return hijack(w)
	})

	if err != nil {
		// If the error path is taken, we never wrapped the connection
		// with the accounting wrapper, so return the resources.
		limits.Remove(ip)
		log.Println(r.RemoteAddr, ":", err)
		if neterr, ok := err.(net.Error); ok && neterr.Timeout() {
			w.WriteHeader(504)
		} else {
			w.WriteHeader(500)
		}
		w.Write([]byte(fmt.Sprintf("%s", err)))
	}
}

var match *utils.Matchmaker
var limits Accounting

func main() {
	flag.Parse()
	limits.connsPerIp = make(map[string]int)
	match = utils.NewMatchmaker()
	match.RendezvousTimeout = time.Duration(*rendezvousTimeout) * time.Second
	match.ProxyTimeout = time.Duration(*inactivityTimeout) * time.Second
	go func() {
		for entry := range match.Log {
			log.Printf("\n%s\n\n", entry)
		}
	}()
	http.HandleFunc("/bounce/", handle)
	http.HandleFunc("/", documentation)
	if err := http.ListenAndServe(":4242", nil); err != nil {
		log.Fatalln(err)
	}
}
